#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>  // 用于固定宽度整数类型（避免溢出）

// 定义大整数类型（根据需求调整，此处用64位整数）
typedef int64_t bigint;

/**
 * 扩展欧几里得算法：计算 ax + by = gcd(a,b)
 * @param a: 方程系数1
 * @param b: 方程系数2
 * @param x: 输出参数，x的解
 * @param y: 输出参数，y的解
 * @param g: 输出参数，gcd(a,b)
 * @return: 成功返回0，内存分配失败返回-1
 */
int extend_euclid(bigint a, bigint b, bigint *x, bigint *y, bigint *g) {
    if (a < b) {
        int ret = extend_euclid(b, a, y, x, g); // 交换x和y
        return ret;
    }

    bigint *coefficients = NULL;
    int coeff_len = 0;
    int coeff_cap = 4;

    // 动态分配系数数组
    coefficients = (bigint *)malloc(coeff_cap * sizeof(bigint));
    if (!coefficients) {
        fprintf(stderr, "扩展欧几里得：内存分配失败！\n");
        return -1;
    }

    // 计算系数 a//b
    while (b != 0) {
        if (coeff_len >= coeff_cap) {
            coeff_cap *= 2;
            bigint *temp = (bigint *)realloc(coefficients, coeff_cap * sizeof(bigint));
            if (!temp) {
                free(coefficients);
                fprintf(stderr, "扩展欧几里得：内存分配失败！\n");
                return -1;
            }
            coefficients = temp;
        }
        coefficients[coeff_len++] = a / b;
        bigint temp = b;
        b = a % b;
        a = temp;
    }

    // 回代计算x, y
    *x = 1;
    *y = 0;
    *g = a;

    for (int i = coeff_len - 1; i >= 0; i--) {
        bigint new_x = *y;
        bigint new_y = *x - (*y) * coefficients[i];
        *x = new_x;
        *y = new_y;
    }

    free(coefficients);
    return 0;
}

/**
 * 丢番图方程求解：ax + by = c
 * @param a: 系数1
 * @param b: 系数2
 * @param c: 常数项
 * @param x0: 输出参数，x特解
 * @param y0: 输出参数，y特解
 * @param xk: 输出参数，x通解系数
 * @param yk: 输出参数，y通解系数
 * @return: 成功返回0，无解返回-1，内存错误返回-2
 */
int diophantine(bigint a, bigint b, bigint c,
                bigint *x0, bigint *y0, bigint *xk, bigint *yk) {
    bigint x_g, y_g, g;
    int ret = extend_euclid(a, b, &x_g, &y_g, &g);
    if (ret != 0) return -2;

    if (c % g != 0) {
        fprintf(stderr, "丢番图方程：无解（%lld 不是 gcd(%lld,%lld)=%lld 的倍数）\n",
                c, a, b, g);
        return -1;
    }

    bigint k = c / g;
    *x0 = x_g * k;
    *y0 = y_g * k;
    *xk = b / g;
    *yk = -a / g;

    return 0;
}

/**
 * Montgomery结构体（对应Python的Montgomery类）
 */
typedef struct {
    bigint r;    // 基数（2的幂）
    bigint n;    // 模数（与r互质）
    bigint r_;   // r在模n下的逆元（r*r_ ≡ 1 mod n）
    bigint n_;   // -n在模r下的逆元（r*r_ - n*n_ = 1）
} Montgomery;

/**
 * 初始化Montgomery结构体
 * @param mont: 待初始化的Montgomery指针
 * @param r: 基数
 * @param n: 模数
 * @return: 成功返回0，失败返回-1
 */
int montgomery_init(Montgomery *mont, bigint r, bigint n) {
    mont->r = r;
    mont->n = n;

    // 求解丢番图方程 r*x + n*y = 1
    bigint x0, y0, xk, yk;
    int ret = diophantine(r, n, 1, &x0, &y0, &xk, &yk);
    if (ret != 0) return -1;

    // 调整r_为正（对应Python的while循环）
    mont->r_ = x0;
    mont->n_ = y0;
    while (mont->r_ < 0) {
        if (xk > 0) {
            mont->r_ += xk;
            mont->n_ += yk;
        } else {
            mont->r_ -= xk;
            mont->n_ -= yk;
        }
    }

    return 0;
}

/**
 * Montgomery约减：将x转换为Montgomery域的约减结果
 * @param mont: Montgomery结构体指针
 * @param x: 待约减的数（x = a*r mod n，a为原数）
 * @return: 约减后结果（∈[0, n-1]）
 */
bigint montgomery_reduce(const Montgomery *mont, bigint x) {
    bigint r = mont->r;
    bigint n = mont->n;
    bigint n_ = mont->n_;

    // 计算 q = (x mod r) * n_ mod r
    bigint q = ((x % r) * n_) % r;
    // 计算 a = (x - q*n) / r
    bigint a = (x - q * n) / r;

    // 确保结果为正
    while (a < 0) {
        a += n;
    }

    return a;
}

/**
 * 将原数x转换为Montgomery域表示（x' = x*r mod n）
 * @param mont: Montgomery结构体指针
 * @param x: 原数
 * @return: Montgomery域表示
 */
bigint montgomery_present(const Montgomery *mont, bigint x) {
    return (x * mont->r) % mont->n;
}

/**
 * 二进制快速幂（Montgomery版本）：计算 (a^e) mod n
 * @param a: 底数
 * @param e: 指数
 * @param n: 模数
 * @return: 结果（∈[0, n-1]）
 */
bigint montgomery_power(bigint a, bigint e, bigint n) {
    // 计算r = 2^k，其中r > n且r与n互质（n为奇数，r是2的幂，自然互质）
    bigint r = 2;
    while (r <= n) {
        r <<= 1; // r *= 2
    }

    // 初始化Montgomery
    Montgomery mont;
    if (montgomery_init(&mont, r, n) != 0) {
        fprintf(stderr, "Montgomery初始化失败！\n");
        return -1;
    }

    // 转换为Montgomery域
    bigint result = montgomery_present(&mont, 1);
    bigint p = montgomery_present(&mont, a);

    bigint power = 1;
    while (power <= e) {
        if ((power & e) != 0) {
            // 累加：result = (result * p) mod n（Montgomery域）
            result = montgomery_reduce(&mont, result * p);
        }
        // 底数平方：p = (p * p) mod n（Montgomery域）
        p = montgomery_reduce(&mont, p * p);
        power <<= 1; // power *= 2
    }

    // 转换回原域
    return montgomery_reduce(&mont, result);
}

// -------------------------- 测试函数 --------------------------
/**
 * 测试Montgomery空间转换
 */
void test_montgomery_space() {
    bigint x = 1024, y = 2048, r = 1024, n = 997;
    // 计算r的模逆（用于验证）
    bigint r_, dummy_y, g;
    extend_euclid(r, n, &r_, &dummy_y, &g);
    while (r_ < 0) r_ += n;

    bigint x_ = (x * r) % n;
    bigint y_ = (y * r) % n;
    bigint xy_verify = (x_ * y_ * r_) % n;

    printf("test_montgomery_space:\n");
    printf("r_ = %lld\n", r_);
    printf("%lld * %lld %% %lld = %lld\n", x, r, n, x_);
    printf("%lld * %lld %% %lld = %lld\n", y, r, n, y_);
    printf("%lld * %lld * %lld %% %lld = %lld\n\n", x_, y_, r_, n, xy_verify);
}

/**
 * 测试Montgomery模幂（核心测试用例）
 */
void test_montgomery_pow() {
    bigint a = 3, e = 777, n = 17;
    bigint r1 = montgomery_power(a, e, n);
    // 直接计算验证（a^e mod n）
    bigint r2 = 1;
    for (bigint i = 0; i < e; i++) {
        r2 = (r2 * a) % n;
    }

    printf("test_montgomery_pow:\n");
    printf("Montgomery结果: %lld, 直接计算结果: %lld\n", r1, r2);
    printf("结果是否一致: %s\n\n", (r1 == r2) ? "True" : "False");
}

/**
 * 测试Montgomery约减
 */
void test_montgomery_reduce() {
    bigint r = 3, n = 4;
    Montgomery mont;
    if (montgomery_init(&mont, r, n) != 0) {
        fprintf(stderr, "test_montgomery_reduce: 初始化失败！\n");
        return;
    }

    bigint x = 1, y = 99;
    bigint x_ = (x * r) % n;
    bigint y_ = (y * r) % n;
    bigint xy_ = (x * y * r) % n;
    bigint reduce_result = montgomery_reduce(&mont, x_ * y_);

    printf("test_montgomery_reduce:\n");
    printf("约减结果: %lld, 预期结果: %lld\n", reduce_result, xy_);
}

// 主函数（执行所有测试）
int main() {
    test_montgomery_space();
    test_montgomery_pow();   // 核心测试用例，对应Python的test_montgomery_pow
    test_montgomery_reduce();
    return 0;
}